################################################################################
class rap_states:
  """ Superclass for storing state data: CA, LS, or IC. """
  def __init__(self):
    self.clear()

# ------------------------------------------------------------------------------

  def __len__(self):
    return self.num_levels()

# ------------------------------------------------------------------------------

  def __getitem__(self,key):
    return self.label(key)

# ------------------------------------------------------------------------------

  def __contains__(self,i):
    return self.v_energy.__contains__(i)

# ------------------------------------------------------------------------------

  def __iter__(self):
    return self.v_energy.__iter__()

# ------------------------------------------------------------------------------

  def _mapme(self,idx,imap=True):
    tidx=idx
    if imap: tidx=self.lmap(tidx)
    if tidx is None: return None
    if tidx not in self: raise ErrorRAP('level index out of range')
    return tidx

# ------------------------------------------------------------------------------

  def _mapme_r(self,idx,imap=True):
    tidx=idx
    if imap: tidx=self.lmap_reverse(tidx)
    if tidx is None: return None
    return tidx

# ------------------------------------------------------------------------------

  def _map_straight(self,other,fromlist,tolist):
    """ 
    General mapping function based on energy order for levels in fromlist and
    tolist grouped by symmetry.
    """
    if type(self) != type(other):
      msg="cannot map different types"
      raise ErrorRAP(msg)

    out={}
    for key,ifrom in fromlist.iteritems():
      if key not in tolist: continue
      ito=tolist[key]

      # sort indices by energy
      tmp=sorted([(self.energy(i),i) for i in ifrom])
      tifrom=[j for (i,j) in tmp]
      tmp=sorted([(other.energy(i),i) for i in ito])
      tito=[j for (i,j) in tmp]

      tmap={}
      while len(tifrom) > 0 and len(tito) > 0: 
        tmap[tito.pop(0)]=tifrom.pop(0)
      for i in ito: 
        if i not in tmap: tmap[i]=None
      for k,v in tmap.iteritems(): out[k]=v
    return out

# ------------------------------------------------------------------------------

  def _map_by_energy(self,other,fromlist,tolist):
    """ 
    General mapping function minimizing energy differences for levels in 
    fromlist and tolist grouped by symmetry.
    """
    if type(self) != type(other):
      msg="cannot map different types"
      raise ErrorRAP(msg)

    out={}
    for key,ifrom in fromlist.iteritems():
      if key not in tolist: continue
      ito=tolist[key]

      # sort indices by energy
      tmp=sorted([(self.energy(i),i) for i in ifrom])
      tifrom=[j for (i,j) in tmp]
      tmp=sorted([(other.energy(i),i) for i in ito])
      tito=[j for (i,j) in tmp]

      tmap={}
      for i in tifrom:
        if len(tito) == 0: continue
        ediff=[abs(self.energy(i)-other.energy(j)) for j in tito]
        jj=ediff.index(min(ediff))
        tmap[tito[jj]]=i
        tito=tito[jj:]
      for i in ito: 
        if i not in tmap: tmap[i]=None
      for k,v in tmap.iteritems(): out[k]=v
    return out

# ------------------------------------------------------------------------------

  def clear(self):
    """ Erase all data. """
    self.v_cup=None            # coupling scheme defined by child type
    self.v_zeff=None
    self.v_nelec=None
    self.v_core_energy=None    # in Ry
    self.v_ion_thresh=None     # in Ry
    self.v_parity={}
    self.v_energy={}           # in Ry
    self.v_lmap={}
    self.v_lmap_r={}           # reverse map

# ------------------------------------------------------------------------------

  def clear_levels(self):
    """ Just erase level data. """
    self.v_parity={}
    self.v_energy={}
    self.v_lmap={}
    self.v_lmap_r={}

# ------------------------------------------------------------------------------

  def copy(self):
    """ Return a copy. """
    return copy.deepcopy(self)

# ------------------------------------------------------------------------------

  def loadfile(self,filename):
    """ Load in level information from level data file. """
    raise 'No input data types defined'

# ------------------------------------------------------------------------------

  def coupling(self):
    """ Return coupling scheme. """
    return self.v_cup

# ------------------------------------------------------------------------------

  def indices(self,imap=False):
    """ Return list of level indices. """
    out=self.v_energy.keys()
    if imap is True: out=map(self._mapme_r,out)
    out.sort()
    return out

# ------------------------------------------------------------------------------

  def reindex(self,remap=None):
    """ 
    Reassign level indices.  Indices are remapped according to dictionary,
    remap, old (key) -> new (value).  If remap is None, then reindex according
    to energy order.
    """
    if remap:
      err=False
      for i in self.indices():
        if i not in remap:
          err=True
          break
      if err:
        raise ErrorRAP("all indices must be represented using reindex")
      imap=remap
    else:
      energies=[]
      for i,e in self.v_energy.iteritems(): energies.append((e,i))
      energies.sort()
      imap={}
      for j in range(len(energies)): imap[energies[j][1]]=j+1
    self.v_parity=dict([(imap[i],val) for (i,val) in self.v_parity.items()])
    self.v_energy=dict([(imap[i],val) for (i,val) in self.v_energy.items()])
    self.v_lmap=dict([(imap[i],val) for (i,val) in self.v_lmap.items()])
    self.v_lmap_r=dict([(v,k) for k,v in self.v_lmap.items()])
    return imap

# ------------------------------------------------------------------------------

  def num_levels(self):
    """ Return number of states. """
    return len(self.v_energy)

# ------------------------------------------------------------------------------

  def nuclear_charge(self):
    """ Return nuclear charge. """
    return self.v_zeff

# ------------------------------------------------------------------------------

  def set_nuclear_charge(self,zeff):
    """ Specify nuclear charge. """
    if zeff is None:
      self.v_zeff=None
    else:
      self.v_zeff=int(zeff)

# ------------------------------------------------------------------------------

  def unset_nuclear_charge(self):
    """ Remove nuclear charge specification. """
    self.v_zeff=None

# ------------------------------------------------------------------------------

  def num_electrons(self):
    """ Return the number of electrons. """
    if self.v_nelec == 0: raise ErrorRAP('no levels present')
    return self.v_nelec

# ------------------------------------------------------------------------------

  def core_energy(self):
    """ Return total energy of all electrons relative to bare nucleus. """
    return self.v_core_energy

# ------------------------------------------------------------------------------

  def set_core_energy(self,core):
    """ Set the core energy in Ry. """
    if core is None:
      self.v_core_energy=None
    else:
      self.v_core_energy=float(core)

# ------------------------------------------------------------------------------

  def unset_core_energy(self):
    """ Remove core energy specification. """
    self.v_core_energy=None

# ------------------------------------------------------------------------------

  def ionization_threshold(self,idx=None):
    """ Return ionization potential of given level (default: ground). """
    if self.v_ion_thresh is None or idx is None: return self.v_ion_thresh
    if idx not in self:
      raise ErrorRAP('level index out of range')
    return self.v_ion_thresh-self.energy(idx)

# ------------------------------------------------------------------------------

  def ionization_potential(self,idx=None):
    """ Return ionization potential of given level (default: ground). """
    return self.ionization_threshold(idx)

# ------------------------------------------------------------------------------

  def set_ionization_threshold(self,ionpot):
    """ Set ionization potential of ground state. """
    if ionpot is None:
      self.v_ion_thresh=None
    else:
      self.v_ion_thresh=float(ionpot)

# ------------------------------------------------------------------------------

  def set_ionization_potential(self,ionpot):
    """ Set ionization potential of ground state. """
    self.set_ionization_threshold(ionpot)

# ------------------------------------------------------------------------------

  def unset_ionization_threshold(self):
    """ Remove ionization potential specification. """
    self.v_ion_thresh=None

# ------------------------------------------------------------------------------

  def parity(self,idx,imap=False):
    """ Return parity of level with given index. """
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return self.v_parity[tidx]

# ------------------------------------------------------------------------------

  def energy_range(self,levlist=None,imap=False):
    """ Return tuple containing the minimum and maximum energies. """
    if imap is True: levlist=map(self._mapme,levlist)
    ens=[self.v_energy[i] for i in levlist if i is not None]
    return (min(ens),max(ens))

# ------------------------------------------------------------------------------

  def energy(self,idx,imap=False):
    """ Return energy of level with given index. """
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return self.v_energy[tidx]

# ------------------------------------------------------------------------------

  def set_energy(self,idx,energy,imap=False):
    """ Set energy of given level to new value. """
    tidx=self._mapme(idx,imap)
    self.v_energy[tidx]=energy

# ------------------------------------------------------------------------------

  def level(self,idx,imap=False):
    """ Return all information for single level (par,E) """
    tidx=self._mapme(idx,imap)
    if tidx is None: return (None,None)
    out=(self.v_parity[tidx], self.v_energy[tidx])
    return out

# ------------------------------------------------------------------------------

  def label(self,idx,imap=False,spacer=None):
    """ Return level label; optional spacer can be used instead of space. """
    raise 'label() is undefined'

# ------------------------------------------------------------------------------

  def weight(self,idx,imap=False):
    """ Return weight of level. """
    raise 'weight() is undefined'

# ------------------------------------------------------------------------------

  def wavelength(self,idx,jdx,imap=False):
    """ Return wavelength in Angstroms calculated from energies. """
    tidx=self._mapme(idx,imap)
    tjdx=self._mapme(jdx,imap)
    if tidx is None or tjdx is None: return None
    ei=self.energy(tidx)
    ej=self.energy(tjdx)
    return HC_RYDANG/abs(ei-ej)

# ------------------------------------------------------------------------------

  def shift_energy(self,en):
    """ Shift all energies by specified amount. """
    ten=float(en)
    for i in self: self.v_energy[i]+=ten

# ------------------------------------------------------------------------------

  def pfilter(self,par,subset=None,imap=False):
    """ Return a list of levels with parity par """
    if par not in [0,1]:
      raise ErrorRAP('invalid parity')
    out=[]
    indices=self.indices()
    if subset is not None: indices=subset
    for i in indices:
      if self.parity(i) == par: out.append(i)
    if imap: out=self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def efilter(self,emin,emax,subset=None,imap=False):
    """ Return a list of levels between emin and emax """
    if emax < emin: emin,emax=emax,emin
    out=[]
    indices=self.indices()
    if subset is not None: indices=subset
    for i in indices:
      e=self.energy(i)
      if e >= emin and e <= emax: out.append(i)
    if imap: out=self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def set_lmap(self,other):
    raise 'set_lmap() undefined'

# ------------------------------------------------------------------------------

  def lmap(self,idx=None):
    """ 
    Return mapped indices.

    If no argument is passed, then return the entire mapping dictionary.  
    Otherwise, the argument can be an integer, list of integers, 2-element
    tuple, or list of 2-element tuples.  The last two items are for mapping
    transitions.
    """
    if idx is None: return self.v_lmap
    if len(self.v_lmap) == 0: return idx
    ttype=type(idx)
    chk=idx
    if ttype == type(0) or ttype == type((1,)): chk=[chk]
    out=[]
    if len(chk) > 0 and isinstance(chk[0],tuple):
      for (i,j) in chk:
        if i in self.v_lmap and j in self.v_lmap:
          out.append((self.v_lmap[i],self.v_lmap[j]))
        else:
          out.append(None)
#      out=[(self.v_lmap[i],self.v_lmap[j]) if i in self.v_lmap and \
#           j in self.v_lmap else None for (i,j) in chk]
    else:
      for i in chk:
        if i in self.v_lmap:
          out.append(self.v_lmap[i])
        else:
          out.append(None)
#      out=[self.v_lmap[i] if i in self.v_lmap else None for i in chk]
    if ttype == type(0) or ttype == type((1,)): out=out[0]
    return out

# ------------------------------------------------------------------------------

  def lmap_reverse(self,idx=None):
    """ 
    Return original indices given mapped values.

    If no argument is passed, then return the entire mapping dictionary.  
    Otherwise, the argument can be an integer, list of integers, 2-element
    tuple, or list of 2-element tuples.  The last two items are for mapping
    transitions.
    """
    if idx is None: return self.v_lmap_r
    if len(self.v_lmap_r) == 0: return idx
    ttype=type(idx)
    chk=idx
    if ttype == type(0) or ttype == type((1,)): chk=[chk]
    out=[]
    if len(chk) > 0 and isinstance(chk[0],tuple):
      for (i,j) in chk:
        if i in self.v_lmap_r and j in self.v_lmap_r:
          out.append((self.v_lmap_r[i],self.v_lmap_r[j]))
        else:
          out.append(None)
#      out=[(self.v_lmap_r[i],self.v_lmap_r[j]) if i in self.v_lmap_r and \
#           j in self.v_lmap_r else None for (i,j) in chk]
    else:
      for i in chk:
        if i in self.v_lmap_r:
          out.append(self.v_lmap_r[i])
        else:
          out.append(None)
#      out=[self.v_lmap_r[i] if i in self.v_lmap_r else None for i in chk]
    if ttype == type(0) or ttype == type((1,)): out=out[0]
    return out

# ------------------------------------------------------------------------------

  def add_level(self):
    raise 'add_level() undefined'

# ------------------------------------------------------------------------------

  def add_levels(self):
    raise 'add_level() undefined'

